# sources/generators.py
import json
import random
import time
from datetime import datetime, timedelta
from confluent_kafka import Producer

class KafkaDataProducer:
    def __init__(self, bootstrap_servers, topic):
        self.conf = {
            'bootstrap.servers': bootstrap_servers,
            'client.id': 'wordcount-producer'
        }
        self.topic = topic
        self.producer = Producer(self.conf)
        self.words = ['apple', 'banana', 'cherry', 'date', 'elderberry']
        self.flag = [True, False]
    
    def _generate_data(self):
        """生成带时间戳的模拟数据"""
        base_time = datetime.now()
        time_offset = random.uniform(-20, 20)
        flag = random.choice(self.flag)
        return {
            'word': random.choice(self.words),
            'count': 1,
            'timestamp': (base_time + timedelta(seconds=time_offset)).isoformat(),
            'is_retract': flag,
        }
    
    def _delivery_report(self, err, msg):
        """消息发送回调函数"""
        if err:
            print(f'消息发送失败: {err}')
        # else:
        #     print(f'消息已送达 [{msg.topic()}] 分区: {msg.partition()}')

    def run(self, interval=1.0):
        """持续生产数据到Kafka"""
        try:
            while True:
                data = self._generate_data()
                print(f"数据：{data}")
                self.producer.produce(
                    self.topic,
                    value=json.dumps(data),
                    callback=self._delivery_report
                )
                self.producer.poll(0)
                time.sleep(interval)
        except KeyboardInterrupt:
            print("\n生产者已停止")
        finally:
            self.producer.flush()